# 实现嵌入计算的函数
def 自定义嵌入层(输入的张量, 嵌入层矩阵):
    """
    torch.nn.embedding
    这里只有前向传播的过程
    :param 输入的张量:
    :param 嵌入层矩阵: 权重矩阵
    :return:
    """
    输出 = []
    for i in range(输入的张量.size(0)):  # 遍历批次
        嵌入样本列表 = []
        for j in range(输入的张量.size(1)):  # 遍历序列中的每个位置
            索引 = 输入的张量[i, j]
            嵌入层向量 = 嵌入层矩阵[索引]
            嵌入样本列表.append(嵌入层向量)
        输出.append(torch.stack(嵌入样本列表))
    return torch.stack(输出)


if __name__ == '__main__':

    """
    在上述代码中：
    首先定义了词汇表大小词汇表大小和嵌入维度嵌入维度，然后随机初始化了一个嵌入层矩阵嵌入层矩阵，它类似于torch.nn.Embedding中的权重矩阵。
    接着模拟生成了一个输入张量输入的张量，其元素是词汇表中的随机索引，形状为(批长, 序长)。
    定义了自定义嵌入层函数来实现嵌入计算。在函数内部，通过两层循环分别遍历批次和序列中的每个位置，根据输入张量中的索引从嵌入层矩阵中取出对应的嵌入向量，最后将每个样本的嵌入向量堆叠起来形成输出张量。
    最后执行自定义的嵌入计算并打印出输入和输出张量的形状，可以看到输出张量的形状为(批长, 序长, 嵌入维度)，符合嵌入计算的预期结果。
    请注意，实际应用中，可能需要根据具体情况对代码进行调整，比如更合理地初始化嵌入层矩阵（例如使用预训练的词向量初始化）、处理不同的数据类型和形状等。而且这里只是一个简单的示例来展示原理，torch.nn.Embedding在内部实现上会更加高效和优化，通常在实际项目中还是推荐直接使用它来进行嵌入操作。
    """
    import torch

    # 假设的词汇表大小和嵌入维度
    词汇表大小 = 10
    嵌入维度 = 5

    # 随机初始化嵌入层矩阵（权重），相当于 nn.Embedding中的权重矩阵
    嵌入层矩阵 = torch.randn(词汇表大小, 嵌入维度)

    # 模拟输入的单词索引张量，形状为 (批长, 序长)
    批长 = 3
    序长 = 4
    输入的张量 = torch.randint(0, 词汇表大小, (批长, 序长))

    # 执行自定义的嵌入计算
    输出的张量 = 自定义嵌入层(输入的张量, 嵌入层矩阵)

    print("输入张量形状:", 输入的张量.shape)
    print("输出张量形状:", 输出的张量.shape)
